package com.sheepone.util.weka;

import weka.clusterers.SimpleKMeans;
import weka.core.Instance;
import weka.core.Instances;

import java.util.List;

/**
 * @author Miss.杨
 * @description
 * @since 2024/3/4 - 17:26
 */
public class KmeansResult {
    private SimpleKMeans kmeans;


    private Instances instances;

    private int[] clusterAssignments;

    private List<Instance> clusterCentroids;


    public KmeansResult(SimpleKMeans kmeans, Instances instances)  {
        this.kmeans = kmeans;
        this.instances = instances;

        // 设置 聚类中心集
        this.clusterCentroids = kmeans.getClusterCentroids();

        // 获取实例对应的聚类标签
        int[] clusterAssignments = new int[this.instances.numInstances()];
        for (int i = 0; i < this.instances.numInstances(); i++) {
            try {
                clusterAssignments[i] = kmeans.clusterInstance(instances.instance(i));
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
        this.clusterAssignments = clusterAssignments;
    }

    public SimpleKMeans getKmeans() {
        return kmeans;
    }

    public int[] getClusterAssignments() {
        return clusterAssignments;
    }

    public List<Instance> getClusterCentroids() {
        return clusterCentroids;
    }

    public Instances getInstances() {
        return instances;
    }
}
